package cn.lg.soar.mvc.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.ByteArrayOutputStream;
import java.io.IOException;

/**
 * 响应包装(支持多次读写)
 * @author luguoxiang
 * @date 2022-04-07
 * 开源项目：https://gitee.com/lgx1992/lg-soar 求star！请给我star！请帮我点个star！
 */
public class SoarResponseWrapper extends HttpServletResponseWrapper {

    protected static final Logger LOGGER = LoggerFactory.getLogger(SoarResponseWrapper.class);

    /**
     * 我们的分支流
     */
    private ByteArrayOutputStream output;

    private ServletOutputStream filterOutput;

    public SoarResponseWrapper(HttpServletResponse response) {
        super(response);
        // 包装输出流，在写入时写我们的字节流
        filterOutput = new ServletOutputStream() {
            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {}

            /**
             * 原理就是往response流中写入时，往我们的ByteArrayOutputStream流中写入
             */
            @Override
            public void write(int b) throws IOException {
                output.write(b);
            }
        };
        output = new ByteArrayOutputStream();
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return filterOutput;
    }

    public byte[] toByteArray() {
        return output.toByteArray();
    }

    public byte[] toByteArrayAndClear() {
        byte[] bytes = output.toByteArray();
        clear();
        return bytes;
    }

    public byte[] toByteArrayAndReset() {
        byte[] bytes = output.toByteArray();
        reset();
        return bytes;
    }

    public void clear() {
        try {
            output.close();
        } catch (IOException e) {
            LOGGER.error("登录响应内存释放失败", e);
        }
    }

    @Override
    public void reset() {
        clear();
        output = new ByteArrayOutputStream();
        super.reset();
    }

}
